
import numpy as np
from utils.numeric import *
from algorithm.source_terms.flow_patterns.source_base import BaseSource
from algorithm.source_terms.source_base_formula import *
from utils.crdlog import log
from utils.constant import *

class StratifiedFlowCalculator(BaseSource):
    def __init__(self, standard_input: dict):
        self.standard_input = standard_input
        self.liquid_height_segment = None
        self.circle_center_angle_segment = None
        self.gas_wet_length_segment = None
        self.liquid_wet_length_segment = None
        self.interface_wet_length_segment = None
        self.gas_hydraulic_diameter_segment = None
        self.Re_g_segment = None
        self.liquid_hydraulic_diameter_segment = None
        self.Re_L_segment = None
        self.rho_gD_segment = None
        self.V_sgt_segment = None
        self.gas_wall_friction_coefficient = None

    def _cal_liquid_height(self, D_segment: float, alpha_L_segment: np.ndarray, rho_L_segment: np.ndarray, rho_g_segment: np.ndarray, V_L_segment: np.ndarray, V_g_segment: np.ndarray, tilt_angle_segment: np.ndarray, liquid_surface_tension_segment: np.ndarray) -> np.ndarray:
        """ 利用持液率求解液位高度

        Args:
            D_segment : float 管的内径
            alpha_L_segment : np.ndarray 液相截面分数

        Returns:
            np.ndarray 液位高度
        """

        self.liquid_height_segment = D_segment * (1 - np.cos(np.pi * alpha_L_segment)) / 2
        
        # # Calculate velocity difference and tilt angle conditions
        # velocity_diff = np.abs(V_g_segment - V_L_segment)
        # small_tilt = np.abs(tilt_angle_segment) < 1e-7
        # small_velocity = velocity_diff < 1e-7

        # # Calculate square root term
        # square_root = (rho_g_segment*velocity_diff**2/2/(rho_L_segment - rho_g_segment)/Physics.g.value/(np.sin(tilt_angle_segment) + 1e-7))**2 - \
        #              4*liquid_surface_tension_segment/(rho_L_segment - rho_g_segment)/Physics.g.value/(np.sin(tilt_angle_segment) + 1e-7)

        # # Calculate height for non-zero square root case
        # height_nonzero = 0.5*(rho_g_segment*velocity_diff**2/2/(rho_L_segment - rho_g_segment)/Physics.g.value/(np.sin(tilt_angle_segment) + 1e-7) + \
        #                      np.sqrt(np.maximum(square_root, 0)))

        # # Calculate height for small tilt, non-zero velocity case
        # height_small_tilt = liquid_surface_tension_segment / (rho_L_segment - rho_g_segment) / Physics.g.value / np.sin(tilt_angle_segment)

        # # Combine conditions
        # self.liquid_height_segment = np.where(small_tilt & small_velocity, 0,
        #                             np.where(small_tilt & ~small_velocity, height_small_tilt,
        #                             np.where(square_root < 0, 0, height_nonzero)))




    def _cal_circle_center_angle(self, liquid_height_segment: np.ndarray, D_segment: float) -> np.ndarray:
        """ 计算液位对应的圆心角

        Args:
            liquid_height_segment : np.ndarray 管内液位高度
            D_segment : float 管的内径

        Returns:
            np.ndarray 圆心角
        """
        self.circle_center_angle_segment = 2 * np.arccos(1 - 2 * liquid_height_segment / D_segment)

    def _cal_liquid_wet_length(self, D_segment: np.ndarray, circle_center_angle: np.ndarray) ->np.ndarray:
        """ 计算液相湿周

            Args:
                D_segment (np.ndarray): 管的内径
                circle_center_angle (np.ndarray): 湿周圆心角

        Returns:
            np.ndarray: 液相的湿周
        """
        self.liquid_wet_length_segment = D_segment * circle_center_angle / 2

    def _cal_gas_wet_length(self, D_segment: np.ndarray, circle_center_angle: np.ndarray) ->np.ndarray:
        """ 计算气相的湿周

        Args:
            D_segment (np.ndarray): 管的内径
            circle_center_angle (np.ndarray): 湿周圆心角

        Returns:
            np.ndarray: 气相湿周
        """
        self.gas_wet_length_segment = D_segment * (np.pi - circle_center_angle / 2)

    def _cal_interface_wet_length(self, D_segment: np.ndarray, circle_center_angle: np.ndarray) ->np.ndarray:
        """ 计算气液界面湿周

        Args:
            D_segment (np.ndarray): 管的内径
            circle_center_angle (np.ndarray): 湿周圆心角

        Returns:
            np.ndarray: 气液界面湿周
        """
        self.interface_wet_length_segment = D_segment * np.sin(circle_center_angle / 2)
    
    def _cal_gas_hydraulic_diameter(self, gas_wet_length: np.ndarray, interface_wet_length: np.ndarray, cross_section_area_segment: np.ndarray, alpha_g_segment: np.ndarray) -> np.ndarray:
        """ 计算气相水力直径

        Args:
            gas_wet_length : np.ndarray 气相湿周
            interface_wet_length : np.ndarray 气液界面湿周
            cross_section_area_segment : np.ndarray 管道截面面积

        Returns:
            np.ndarray 气相水力直径
        """
        self.gas_hydraulic_diameter_segment = 4 * alpha_g_segment * cross_section_area_segment / (gas_wet_length + interface_wet_length)
    
    def _cal_gas_Reynolds_number(self, rho_g_segment: np.ndarray, gas_hydraulic_diameter_segment: np.ndarray, V_g_segment: np.ndarray, gas_viscosity_segment: np.ndarray) -> np.ndarray:
        """ 计算气相雷诺数

        Args:
            rho_g_segment : np.ndarray 气相密度
            gas_hydraulic_diameter_segment : np.ndarray 水力直径
            V_g_segment : np.ndarray 气相速度
            gas_viscosity_segment : np.ndarray 气相粘度

        Returns:
            np.ndarray 气相雷诺数
        """
        self.Re_g_segment = cal_Reynolds_number(rho_g_segment, gas_hydraulic_diameter_segment, V_g_segment, gas_viscosity_segment)
    
    def _cal_liquid_hydraulic_diameter(self, liquid_wet_length: np.ndarray, gas_wet_length: np.ndarray, cross_section_area_segment: np.ndarray, alpha_L_segment: np.ndarray) -> np.ndarray:
        """ 计算液相水力直径
        
        Args:
            liquid_wet_length : np.ndarray 液相湿周
            cross_section_area_segment : np.ndarray 管道截面面积
            alpha_L_segment : np.ndarray 液相截面分数

        Returns:
            np.ndarray 液相水力直径
        """
        self.liquid_hydraulic_diameter_segment = 4 * alpha_L_segment * cross_section_area_segment / (liquid_wet_length + gas_wet_length)
    
    def _cal_liquid_Reynolds_number(self, rho_L_segment: np.ndarray, liquid_hydraulic_diameter_segment: np.ndarray, V_L_segment: np.ndarray, liquid_viscosity_segment: np.ndarray) -> np.ndarray:
        """ 计算液相雷诺数
        """
        self.Re_L_segment = cal_Reynolds_number(rho_L_segment, liquid_hydraulic_diameter_segment, V_L_segment, liquid_viscosity_segment)
    
    @staticmethod
    def _cal_wall_friction_coefficient(Re_segment: np.ndarray, hydraulic_diameter_segment: float, wall_roughness_segment: float) -> np.ndarray:
        """ 气液相和壁面的摩擦系数统一计算方法
        
        Args:
            Re_segment : np.ndarray 雷诺数
            hydraulic_diameter_segment : float 水力直径
            wall_roughness_segment : float 介质表面粗糙度

        Returns:
            np.ndarray 壁面摩擦系数
        """
        friction_coefficient_laminar = np.where(Re_segment > 0, 64/(Re_segment + 1e-12), 0)
        
        friction_coefficient_turbulent = 0.0055 * (1 + (np.where(hydraulic_diameter_segment > 0, (2e4 * wall_roughness_segment) / (hydraulic_diameter_segment + 1e-12), 0) + np.where(Re_segment > 0, 1e6 / (Re_segment + 1e-12), 0)) ** (1 / 3))
        
        friction_coefficient_transition = np.where(friction_coefficient_laminar >= friction_coefficient_turbulent, friction_coefficient_laminar, friction_coefficient_turbulent)
        conditions = [Re_segment < 2000, (Re_segment >= 2000) & (Re_segment <= 3000), Re_segment > 3000]
        choices = [friction_coefficient_laminar, friction_coefficient_transition, friction_coefficient_turbulent]
        wall_friction_coefficient = np.select(conditions, choices)
        return wall_friction_coefficient

    def cal_gas_wall_friction_coefficient(self):
        """ 计算气相和壁面的摩擦系数
        
        """
        wall_roughness_segment = self.standard_input['wall_roughness_segment']
        rho_g_segment = self.standard_input['rho_g_segment']
        cross_section_area_segment = self.standard_input['cross_section_area_segment']
        alpha_g_segment = self.standard_input['alpha_g_segment']
        alpha_L_segment = self.standard_input['alpha_L_segment']
        V_g_segment = self.standard_input['V_g_segment']
        gas_viscosity_segment = self.standard_input['gas_viscosity_segment']
        D_segment = self.standard_input['D_segment']
        rho_L_segment = self.standard_input['rho_L_segment']
        V_L_segment = self.standard_input['V_L_segment']
        tilt_angle_segment = self.standard_input['tilt_angle_segment']
        liquid_surface_tension_segment = self.standard_input['liquid_surface_tension_segment']
        # 开始计算
        if self.liquid_height_segment is None:
            self._cal_liquid_height(D_segment, alpha_L_segment, rho_L_segment, rho_g_segment, V_L_segment, V_g_segment, tilt_angle_segment, liquid_surface_tension_segment)
        if self.circle_center_angle_segment is None:
            self._cal_circle_center_angle(self.liquid_height_segment, D_segment)
        if self.gas_wet_length_segment is None:

            self._cal_gas_wet_length(D_segment, self.circle_center_angle_segment)
        if self.interface_wet_length_segment is None:
            self._cal_interface_wet_length(D_segment, self.circle_center_angle_segment)
        if self.gas_hydraulic_diameter_segment is None:
            self._cal_gas_hydraulic_diameter(self.gas_wet_length_segment, self.interface_wet_length_segment, cross_section_area_segment, alpha_g_segment)
        if self.Re_g_segment is None:
            self._cal_gas_Reynolds_number(rho_g_segment, self.gas_hydraulic_diameter_segment, V_g_segment, gas_viscosity_segment)
        
        self.gas_wall_friction_coefficient = self._cal_wall_friction_coefficient(self.Re_g_segment, self.gas_hydraulic_diameter_segment, wall_roughness_segment)
        return self.gas_wall_friction_coefficient

    def cal_liquid_wall_friction_coefficient(self):
        """ 计算液相和壁面的摩擦系数
        
        """
        rho_L_segment = self.standard_input['rho_L_segment']
        V_L_segment = self.standard_input['V_L_segment']
        liquid_viscosity_segment = self.standard_input['liquid_viscosity_segment']
        D_segment = self.standard_input['D_segment']
        alpha_L_segment = self.standard_input['alpha_L_segment']
        cross_section_area_segment = self.standard_input['cross_section_area_segment']
        wall_roughness_segment = self.standard_input['wall_roughness_segment']
        rho_g_segment = self.standard_input['rho_g_segment']
        V_g_segment = self.standard_input['V_g_segment']
        tilt_angle_segment = self.standard_input['tilt_angle_segment']
        liquid_surface_tension_segment = self.standard_input['liquid_surface_tension_segment']

        # 开始计算
        if self.liquid_height_segment is None:  
            self._cal_liquid_height(D_segment, alpha_L_segment, rho_L_segment, rho_g_segment, V_L_segment, V_g_segment, tilt_angle_segment, liquid_surface_tension_segment)
        if self.gas_wet_length_segment is None:
            self._cal_gas_wet_length(D_segment, self.circle_center_angle_segment)
        if self.circle_center_angle_segment is None:
            self._cal_circle_center_angle(self.liquid_height_segment, D_segment)
        if self.liquid_wet_length_segment is None:
            self._cal_liquid_wet_length(D_segment, self.circle_center_angle_segment)
        if self.liquid_hydraulic_diameter_segment is None:
            self._cal_liquid_hydraulic_diameter(self.liquid_wet_length_segment, self.gas_wet_length_segment, cross_section_area_segment, alpha_L_segment)
        if self.Re_L_segment is None:
            self._cal_liquid_Reynolds_number(rho_L_segment, self.liquid_hydraulic_diameter_segment, V_L_segment, liquid_viscosity_segment)
        liquid_wall_friction_coefficient = self._cal_wall_friction_coefficient(self.Re_L_segment, self.liquid_hydraulic_diameter_segment, wall_roughness_segment)
        return liquid_wall_friction_coefficient
    
    def _cal_gas_threshold_superficial_velocity(self, p_segment: np.ndarray) -> np.ndarray:
        """ 计算气相临界表观速度
        
        Args:
            p_segment : np.ndarray 压力

        Returns:
            np.ndarray 气相临界表观速度
        """
        self.V_sgt_segment = 5 * (101325 / p_segment) ** 0.5
        
        
    def cal_gas_liquid_interface_friction_coefficient(self):
        """计算气液界面摩擦系数
        
        """
        p_segment = self.standard_input['p_segment']
        alpha_L_segment = self.standard_input['alpha_L_segment']
        D_segment = self.standard_input['D_segment']
        V_sg_segment = self.standard_input['V_sg_segment']
        rho_L_segment = self.standard_input['rho_L_segment']
        rho_g_segment = self.standard_input['rho_g_segment']
        V_L_segment = self.standard_input['V_L_segment']
        V_g_segment = self.standard_input['V_g_segment']
        tilt_angle_segment = self.standard_input['tilt_angle_segment']
        liquid_surface_tension_segment = self.standard_input['liquid_surface_tension_segment']



        # 开始计算
        if self.V_sgt_segment is None:
            self._cal_gas_threshold_superficial_velocity(p_segment)
        
        if self.liquid_height_segment is None:
            self._cal_liquid_height(D_segment, alpha_L_segment, rho_L_segment, rho_g_segment, V_L_segment, V_g_segment, tilt_angle_segment, liquid_surface_tension_segment)


        if self.gas_wall_friction_coefficient is None:
            self.cal_gas_wall_friction_coefficient()
        
        conditions = [V_sg_segment <= self.V_sgt_segment, V_sg_segment > self.V_sgt_segment]
        choices = [self.gas_wall_friction_coefficient, self.gas_wall_friction_coefficient*(1 + 15 * ((self.liquid_height_segment / D_segment) ** 0.5) * (V_sg_segment / self.V_sgt_segment - 1))]
        
        gas_liquid_interface_friction_coefficient = np.select(conditions, choices)
        return gas_liquid_interface_friction_coefficient
    
    def _cal_gas_droplet_mixed_density( 
        self,
        rho_g_segment: np.ndarray,
        rho_L_segment: np.ndarray,
        alpha_g_segment: np.ndarray,
        alpha_D_segment: np.ndarray,
    ) -> np.ndarray:
        """ 计算气和液滴的混合密度
        
        Args:
            rho_g_segment : np.ndarray 气相密度
            rho_L_segment : np.ndarray 液相密度
            alpha_g_segment : np.ndarray 气相截面分数
            alpha_D_segment : np.ndarray 液滴截面分数

        Returns:
            np.ndarray 气和液滴的混合密度
        """
        self.rho_gD_segment = rho_g_segment * alpha_g_segment/ (alpha_g_segment + alpha_D_segment + 1e-12) + rho_L_segment * alpha_D_segment/(alpha_g_segment + alpha_D_segment + 1e-12)

    @staticmethod
    def _cal_friction_force(friction_coefficient: np.ndarray, rho_segment: np.ndarray, V_segment: np.ndarray, wet_length: np.ndarray):
        """ 计算摩擦力统一形式
        
        Args:
            friction_coefficient : np.ndarray 摩擦系数
            rho_segment : np.ndarray 密度
            V_segment : np.ndarray 速度
            wet_length : np.ndarray 湿周

        Returns:
            np.ndarray 摩擦力
        """
        return 1/2*friction_coefficient * rho_segment * np.abs(V_segment) * V_segment * wet_length/4 
    
    def cal_gas_wall_friction_force(self):
        """ 计算气壁摩擦力
        
        """
        rho_g_segment = self.standard_input['rho_g_segment']
        V_g_segment = self.standard_input['V_g_segment']
        D_segment = self.standard_input['D_segment']
        alpha_L_segment = self.standard_input['alpha_L_segment']
        alpha_g_segment = self.standard_input['alpha_g_segment']
        alpha_D_segment = self.standard_input['alpha_D_segment']
        rho_L_segment = self.standard_input['rho_L_segment']
        V_L_segment = self.standard_input['V_L_segment']
        V_g_segment = self.standard_input['V_g_segment']
        tilt_angle_segment = self.standard_input['tilt_angle_segment']
        liquid_surface_tension_segment = self.standard_input['liquid_surface_tension_segment']

        # 开始计算
        if self.rho_gD_segment is None:
            self._cal_gas_droplet_mixed_density(rho_g_segment, rho_L_segment, alpha_g_segment, alpha_D_segment)
       
        if self.liquid_height_segment is None:
            self._cal_liquid_height(D_segment, alpha_L_segment, rho_L_segment, rho_g_segment, V_L_segment, V_g_segment, tilt_angle_segment, liquid_surface_tension_segment)
        if self.circle_center_angle_segment is None:

            self._cal_circle_center_angle(self.liquid_height_segment, D_segment)
        if self.gas_wet_length_segment is None:
            self._cal_gas_wet_length(D_segment, self.circle_center_angle_segment)
        
        if self.gas_wall_friction_coefficient is None:
            self.cal_gas_wall_friction_coefficient()
            
        gas_wall_friction_force = self._cal_friction_force(self.gas_wall_friction_coefficient, self.rho_gD_segment, V_g_segment, self.gas_wet_length_segment)
        return gas_wall_friction_force

    def cal_liquid_wall_friction_force(self):
        """ 计算液壁摩擦力
        
        """
        rho_L_segment = self.standard_input['rho_L_segment']
        V_L_segment = self.standard_input['V_L_segment']
        D_segment = self.standard_input['D_segment']
        alpha_L_segment = self.standard_input['alpha_L_segment']
        rho_g_segment = self.standard_input['rho_g_segment']
        V_g_segment = self.standard_input['V_g_segment']
        tilt_angle_segment = self.standard_input['tilt_angle_segment']
        liquid_surface_tension_segment = self.standard_input['liquid_surface_tension_segment']

        # 开始计算
        if self.liquid_height_segment is None:
            self._cal_liquid_height(D_segment, alpha_L_segment, rho_L_segment, rho_g_segment, V_L_segment, V_g_segment, tilt_angle_segment, liquid_surface_tension_segment)
        if self.circle_center_angle_segment is None:

            self._cal_circle_center_angle(self.liquid_height_segment, D_segment)
        if self.liquid_wet_length_segment is None:
            self._cal_liquid_wet_length(D_segment, self.circle_center_angle_segment)
            
        liquid_wall_friction_coefficient = self.cal_liquid_wall_friction_coefficient()
        liquid_wall_friction_force = self._cal_friction_force(liquid_wall_friction_coefficient, rho_L_segment, V_L_segment, self.liquid_wet_length_segment)
        return liquid_wall_friction_force
    
    def cal_gas_liquid_interface_friction_force(self):
        """ 计算气液界面摩擦力
        
        """
        rho_g_segment = self.standard_input['rho_g_segment']
        rho_L_segment = self.standard_input['rho_L_segment']
        alpha_g_segment = self.standard_input['alpha_g_segment']
        alpha_D_segment = self.standard_input['alpha_D_segment']
        D_segment = self.standard_input['D_segment']
        alpha_L_segment = self.standard_input['alpha_L_segment']
        V_r_segment = self.standard_input['V_r_segment']
        V_L_segment = self.standard_input['V_L_segment']
        V_g_segment = self.standard_input['V_g_segment']
        tilt_angle_segment = self.standard_input['tilt_angle_segment']
        liquid_surface_tension_segment = self.standard_input['liquid_surface_tension_segment']

        # 开始计算
        if self.liquid_height_segment is None:
            self._cal_liquid_height(D_segment, alpha_L_segment, rho_L_segment, rho_g_segment, V_L_segment, V_g_segment, tilt_angle_segment, liquid_surface_tension_segment)
        if self.circle_center_angle_segment is None:

            self._cal_circle_center_angle(self.liquid_height_segment, D_segment)
        if self.interface_wet_length_segment is None:   
            self._cal_interface_wet_length(D_segment, self.circle_center_angle_segment)
        if self.rho_gD_segment is None:
            self._cal_gas_droplet_mixed_density(rho_g_segment, rho_L_segment, alpha_g_segment, alpha_D_segment)
        
        gas_liquid_interface_friction_coefficient = self.cal_gas_liquid_interface_friction_coefficient()
        gas_liquid_interface_friction_force = self._cal_friction_force(gas_liquid_interface_friction_coefficient, self.rho_gD_segment, V_r_segment, self.interface_wet_length_segment)
        return gas_liquid_interface_friction_force

